/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

package com.ming.tflitesingleblur.segment;

import android.app.Activity;
import android.graphics.Bitmap;
import android.graphics.Color;
import android.graphics.RectF;
import android.os.SystemClock;
import android.os.Trace;
import android.util.Log;

import com.ming.tflitesingleblur.segment.operator.Rot90Operator;
import com.ming.tflitesingleblur.utils.ImageSaveUtils;
import com.ming.tflitesingleblur.utils.Logger;

import org.tensorflow.lite.DataType;
import org.tensorflow.lite.Interpreter;
import org.tensorflow.lite.gpu.GpuDelegate;
import org.tensorflow.lite.nnapi.NnApiDelegate;
import org.tensorflow.lite.support.common.FileUtil;
import org.tensorflow.lite.support.common.TensorOperator;
import org.tensorflow.lite.support.common.TensorProcessor;
import org.tensorflow.lite.support.image.ImageProcessor;
import org.tensorflow.lite.support.image.TensorImage;
import org.tensorflow.lite.support.image.ops.ResizeOp;
import org.tensorflow.lite.support.image.ops.ResizeOp.ResizeMethod;
import org.tensorflow.lite.support.label.TensorLabel;
import org.tensorflow.lite.support.tensorbuffer.TensorBuffer;

import java.io.IOException;
import java.nio.MappedByteBuffer;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Comparator;
import java.util.List;
import java.util.Map;
import java.util.PriorityQueue;

/**
 * A classifier specialized to label images using TensorFlow Lite.
 */
public abstract class Classifier {
    private static final Logger LOGGER = new Logger();
    private static Model mCurrentModel;

    /**
     * The model type used for classification.
     */
    public enum Model {

        DOG_CAT,
        HUMAN_DANCE
    }

    /**
     * The runtime device type used for executing classification.
     */
    public enum Device {
        CPU,
        NNAPI,
        GPU
    }

    /**
     * Number of results to show in the UI.
     */
    private static final int MAX_RESULTS = 6;

    /**
     * The loaded TensorFlow Lite model.
     */
    private MappedByteBuffer tfliteModel;

    /**
     * Image size along the x axis.
     */
    private final int imageHeight;

    /**
     * Image size along the y axis.
     */
    private final int imageWidth;

    /**
     * Optional GPU delegate for accleration.
     */
    private GpuDelegate gpuDelegate = null;

    /**
     * Optional NNAPI delegate for accleration.
     */
    private NnApiDelegate nnApiDelegate = null;

    /**
     * An instance of the driver class to run model inference with Tensorflow Lite.
     */
    protected Interpreter tflite;

    /**
     * Options for configuring the Interpreter.
     */
    private final Interpreter.Options tfliteOptions = new Interpreter.Options();

    /**
     * Labels corresponding to the output of the vision model.
     */
    private List<String> labels;

    /**
     * Input image TensorBuffer.
     */
    private TensorImage inputImageBuffer;

    /**
     * Output probability TensorBuffer.
     */
    private final TensorBuffer outputBuffer;

    /**
     * Processer to apply post processing of the output probability.
     */
    private final TensorProcessor probabilityProcessor;

    //模型的输入shape和输出shape
//    {1, height, width, 3}
    private int[] inputShape;
    private int[] outputShape;

    public int[] getOutputShape() {
        //把[1, 224, 224, 3]格式更改成[224, 224, 3]
        int[] outShape = new int[outputShape.length - 1];
        System.arraycopy(outputShape,1, outShape, 0, outShape.length);
        return outShape;
    }
    public int getOutputWidth() {
        return outputShape[2];
    }
    public int getOutputHeight() {
        return outputShape[1];
    }


    /**
     * Creates a classifier with the provided configuration.
     *
     * @param activity   The current Activity.
     * @param model      The model to use for classification.
     * @param device     The device to use for classification.
     * @param numThreads The number of threads to use for classification.
     * @return A classifier with the desired configuration.
     */
    public static Classifier create(Activity activity, Model model, Device device, int numThreads)
            throws IOException {
        mCurrentModel = model;
        switch (model) {

            case DOG_CAT:
                return new DogCatClassifier(activity, device, numThreads);
            case HUMAN_DANCE:
                return new HumanDanceClassifier(activity, device, numThreads);

            default:
                return new DogCatClassifier(activity, device, numThreads);
        }
    }

    /**
     * An immutable result returned by a Classifier describing what was recognized.
     */
    public static class Recognition {
        /**
         * A unique identifier for what has been recognized. Specific to the class, not the instance of
         * the object.
         */
        private final String id;

        /**
         * Display name for the recognition.
         */
        private final String title;

        /**
         * A sortable score for how good the recognition is relative to others. Higher should be better.
         */
        private final Float confidence;

        /**
         * Optional location within the source image for the location of the recognized object.
         */
        private RectF location;

        public Recognition(
                final String id, final String title, final Float confidence, final RectF location) {
            this.id = id;
            this.title = title;
            this.confidence = confidence;
            this.location = location;
        }

        public String getId() {
            return id;
        }

        public String getTitle() {
            return title;
        }

        public Float getConfidence() {
            return confidence;
        }

        public RectF getLocation() {
            return new RectF(location);
        }

        public void setLocation(RectF location) {
            this.location = location;
        }

        @Override
        public String toString() {
            String resultString = "";
            if (id != null) {
                resultString += "[" + id + "] ";
            }

            if (title != null) {
                resultString += title + " ";
            }

            if (confidence != null) {
                resultString += String.format("(%.1f%%) ", confidence * 100.0f);
            }

            if (location != null) {
                resultString += location + " ";
            }

            return resultString.trim();
        }
    }

    /**
     * Initializes a {@code Classifier}.
     */
    protected Classifier(Activity activity, Device device, int numThreads) throws IOException {
        tfliteModel = FileUtil.loadMappedFile(activity, getModelPath());
        switch (device) {
            case NNAPI:
                nnApiDelegate = new NnApiDelegate();
                tfliteOptions.addDelegate(nnApiDelegate);
                break;
            case GPU:
                gpuDelegate = new GpuDelegate();
                tfliteOptions.addDelegate(gpuDelegate);
                break;
            case CPU:
                break;
        }
        tfliteOptions.setNumThreads(numThreads);
        //这里是根据模型生成的模型解释器
        tflite = new Interpreter(tfliteModel, tfliteOptions);

        // Loads labels out from the label file.
        labels = FileUtil.loadLabels(activity, getLabelPath());

        // Reads type and shape of input and output tensors, respectively.
        int imageTensorIndex = 0;
        //在这里加载模型里的输入的尺寸

        // {1, height, width, 3}
        inputShape = tflite.getInputTensor(imageTensorIndex).shape();
        Log.i(TAG, "Classifier: inputImageShape = " + Arrays.toString(inputShape));
        imageHeight = inputShape[1];
        imageWidth = inputShape[2];
        DataType imageDataType = tflite.getInputTensor(imageTensorIndex).dataType();
        int outputTensorIndex = 0;

        //获取模型的输出尺寸，不同模型各有不同
        outputShape = tflite.getOutputTensor(outputTensorIndex).shape();
        Log.i(TAG, "Classifier: outputShape = " + Arrays.toString(outputShape));
        DataType outputTensorDataType = tflite.getOutputTensor(outputTensorIndex).dataType();
        Log.i(TAG, "Classifier: outputTensorDataType" + outputTensorDataType.name());
        // Creates the input tensor. 创建输入张量
        inputImageBuffer = new TensorImage(imageDataType);

        // Creates the output tensor and its processor. 创建输出张量和它的处理器
        outputBuffer = TensorBuffer.createFixedSize(outputShape, outputTensorDataType);

        // Creates the post processor for the output probability.
        probabilityProcessor = new TensorProcessor.Builder().add(getPostprocessNormalizeOp()).build();

        LOGGER.d("Created a Tensorflow Lite Image Classifier.");
    }

    private String TAG = "Classifier";

    int saveIndex = 0;

    /**
     * Runs inference and returns the classification results.
     */
    public List<Recognition> recognizeImage(final Bitmap bitmap, int sensorOrientation) {
        // Logs this method so that it can be analyzed with systrace.
        Trace.beginSection("recognizeImage");

        Trace.beginSection("loadImage");
        long startTimeForLoadImage = SystemClock.uptimeMillis();
        //原本是640*480的图，会经过下面的处理
        inputImageBuffer = loadImage(bitmap, sensorOrientation);
        long endTimeForLoadImage = SystemClock.uptimeMillis();
        Trace.endSection();
        LOGGER.v("Timecost to load the image: " + (endTimeForLoadImage - startTimeForLoadImage));

        // Runs the inference call.
        Trace.beginSection("runInference");
        long startTimeForReference = SystemClock.uptimeMillis();


        //手动处理输入前和输出的数据，模型为猫狗分类
/*        testForLocalCatImageRecogonise();
        Map<String, Float> labeledProbability = new LinkedHashMap<>();
        labeledProbability.put("cat", 0.5f);*/

        Log.i(TAG, "recognizeImage: input" +
                Arrays.toString(inputImageBuffer.getTensorBuffer().getFloatArray()));
        Log.i(TAG, "recognizeImage: input image length = " + inputImageBuffer.getTensorBuffer().getFloatArray().length);

        Log.i(TAG, "recognizeImage: input shape " + Arrays.toString(inputImageBuffer.getTensorBuffer().getShape()));

        TensorBuffer buffer = null;

        tflite.run(inputImageBuffer.getBuffer(), outputBuffer.getBuffer().rewind());


        //排个序，看一下最大的结果的条目
        float[] resultArrs = outputBuffer.getFloatArray();

        int maxValueIndex = -1;
        float maxValue = -10000000000f;
        for (int x = 0; x < resultArrs.length; x++){
            if (resultArrs[x] > maxValue) {
                maxValue = resultArrs[x];
                maxValueIndex = x;
            }
        }
        Log.i(TAG, "recognizeImage: maxValue = "+maxValue + ",index is" + maxValueIndex);

        //不频繁存图
        saveIndex++;
        if (false && saveIndex % 20 == 0) {
            //保存图片，看看输入进算法的，它到底是个什么玩意
            if(buffer != null) {
                int shape[] = buffer.getShape();
                Log.i(TAG, "recognizeImage: inputshpae" + Arrays.toString(shape));
                ImageSaveUtils.saveImage("/sdcard/workspace",
                        "image_" + System.currentTimeMillis() + ".png",
                        buffer.getFloatArray(),
                        shape,
                        true, new float[]{0, 255f});
            }

            int shape[] = inputImageBuffer.getTensorBuffer().getShape();
            Log.i(TAG, "recognizeImage: inputshpae" + Arrays.toString(shape));
            ImageSaveUtils.saveImage("/sdcard/workspace",
                    "image_input_" + System.currentTimeMillis() + ".png",
                    inputImageBuffer.getTensorBuffer().getFloatArray(),
                    shape,
                    true, new float[]{0, 255f});

        }

        long endTimeForReference = SystemClock.uptimeMillis();
        Trace.endSection();
        LOGGER.v("Timecost to run model inference: " + (endTimeForReference - startTimeForReference));


        // Gets the map of label and probability.
        Map<String, Float> labeledProbability =
                new TensorLabel(labels, probabilityProcessor.process(outputBuffer))
                        .getMapWithFloatValue();


        Trace.endSection();


        // Gets top-k results.
        return getTopKProbability(labeledProbability);
    }

    /**
     * Runs inference and returns the classification results.
     */
    public float[] segment(final Bitmap bitmap, int sensorOrientation) {
        // Logs this method so that it can be analyzed with systrace.
        Trace.beginSection("recognizeImage");

        Trace.beginSection("loadImage");
        long startTimeForLoadImage = SystemClock.uptimeMillis();
        //原本是640*480的图，会经过下面的处理
        inputImageBuffer = loadImage(bitmap, sensorOrientation);
        long endTimeForLoadImage = SystemClock.uptimeMillis();
        Trace.endSection();
        LOGGER.v("Timecost to load the image: " + (endTimeForLoadImage - startTimeForLoadImage));

        // Runs the inference call.
        Trace.beginSection("runInference");
        long startTimeForReference = SystemClock.uptimeMillis();


        //手动处理输入前和输出的数据，模型为猫狗分类
/*        testForLocalCatImageRecogonise();
        Map<String, Float> labeledProbability = new LinkedHashMap<>();
        labeledProbability.put("cat", 0.5f);*/

//        Log.i(TAG, "recognizeImage: input" +
//                Arrays.toString(inputImageBuffer.getTensorBuffer().getFloatArray()));
//        Log.i(TAG, "recognizeImage: input image length = " + inputImageBuffer.getTensorBuffer().getFloatArray().length);

        Log.i(TAG, "recognizeImage: input shape " + Arrays.toString(inputImageBuffer.getTensorBuffer().getShape()));

        TensorBuffer buffer = null;

        tflite.run(inputImageBuffer.getBuffer(), outputBuffer.getBuffer().rewind());


        //如果没什么问题，这个就是mask了，实际上还需要转化，这个得根据实际的模型输出内容来看
        float[] resultArrs = outputBuffer.getFloatArray();

        long endTimeForReference = SystemClock.uptimeMillis();
        Trace.endSection();
        LOGGER.v("Timecost to run model inference: " + (endTimeForReference - startTimeForReference));

        return resultArrs;
    }

    /**
     * 反转数组
     * @param array
     * @return
     */
    public static float[] reverseArray(float[] array){
        float [] newArray = new float[array.length];
        for(int i=0; i<newArray.length; i++){
            newArray[i] = array[array.length - i - 1];
        }
        return newArray;
    }

    /**
     * Get 28x28 pixel data for tensorflow input.
     */
    public float[][][][] getPixelData(Bitmap bitmap) {
        if (bitmap == null) {
            return null;
        }
        int width = bitmap.getWidth();
        int height = bitmap.getHeight();
        // Get 28x28 pixel data from bitmap
        int[] pixels = new int[width * height];
        bitmap.getPixels(pixels, 0, width, 0, 0, width, height);

        float[][][][] retPixels = new float[1][height][width][3];
        for (int i = 0; i < pixels.length; ++i) {
            int pix = pixels[i];
            int r = Color.red(pix);
            int g = Color.green(pix);
            int b = Color.blue(pix);
            retPixels[0][i / width][i % width][0] = r;
            retPixels[0][i / width][i % width][1] = g;
            retPixels[0][i / width][i % width][2] = b;

            // Set 0 for white and 255 for black pixel
           /* int pix = pixels[i];
            int b = pix & 0xff;
            retPixels[i] = 0xff - b;*/
        }
        return retPixels;
    }

    /**
     * Closes the interpreter and model to release resources.
     */
    public void close() {
        if (tflite != null) {
            tflite.close();
            tflite = null;
        }
        if (gpuDelegate != null) {
            gpuDelegate.close();
            gpuDelegate = null;
        }
        if (nnApiDelegate != null) {
            nnApiDelegate.close();
            nnApiDelegate = null;
        }
        tfliteModel = null;
    }

    /**
     * Get the image size along the x axis.
     */
    public int getImageHeight() {
        return imageHeight;
    }

    /**
     * Get the image size along the y axis.
     */
    public int getImageWidth() {
        return imageWidth;
    }

    /**
     * Loads input image, and applies preprocessing.
     */
    private TensorImage loadImage(final Bitmap bitmap, int sensorOrientation) {
        // Loads bitmap into a TensorImage.
        inputImageBuffer.load(bitmap);

        //这块是内部检查shape的逻辑，如果为false,shape就不能用
//        int[]shape = inputImageBuffer.getTensorBuffer().getShape();
//        boolean b = (shape.length == 3 || shape.length == 4 && shape[0] == 1)//长度可以是3，4,如果是4,第一个参数必须是1
//                && shape[shape.length - 3] > 0
//                && shape[shape.length - 2] > 0
//                && shape[shape.length - 1] == 3;//要求第三个参数必须是3

        // Creates processor for the TensorImage.
        int cropSize = Math.min(bitmap.getWidth(), bitmap.getHeight());
        int numRoration = sensorOrientation / 90;
        Log.i(TAG, "loadImage: numRoration" + numRoration);//1
        // TODO(b/143564309): Fuse ops inside ImageProcessor.
        //这个看起来是处理图片东西
        ImageProcessor imageProcessor;
        ImageProcessor.Builder builder = new ImageProcessor.Builder()
//                .add(new ResizeWithCropOrPadOp(cropSize, cropSize)) //根据宽度裁
                .add(new ResizeOp(imageHeight, imageWidth, ResizeMethod.BILINEAR))
                .add(new Rot90Operator(numRoration)) //右转90度
                .add(getPreprocessNormalizeOp());
        imageProcessor = builder.build();
        //这里的process实质上是按顺序执行上面设定的Op
        return imageProcessor.process(inputImageBuffer);
    }

    /**
     * Gets the top-k results.
     */
    private static List<Recognition> getTopKProbability(Map<String, Float> labelProb) {
        // Find the best classifications.
        PriorityQueue<Recognition> pq =
                new PriorityQueue<>(
                        MAX_RESULTS,
                        new Comparator<Recognition>() {
                            @Override
                            public int compare(Recognition lhs, Recognition rhs) {
                                // Intentionally reversed to put high confidence at the head of the queue.
                                return Float.compare(rhs.getConfidence(), lhs.getConfidence());
                            }
                        });

        for (Map.Entry<String, Float> entry : labelProb.entrySet()) {
            pq.add(new Recognition("" + entry.getKey(), entry.getKey(), entry.getValue(), null));
        }

        final ArrayList<Recognition> recognitions = new ArrayList<>();
        int recognitionsSize = Math.min(pq.size(), MAX_RESULTS);
        for (int i = 0; i < recognitionsSize; ++i) {
            recognitions.add(pq.poll());
        }
        return recognitions;
    }

    /**
     * Gets the name of the model file stored in Assets.
     */
    protected abstract String getModelPath();

    /**
     * Gets the name of the label file stored in Assets.
     */
    protected abstract String getLabelPath();

    /**
     * Gets the TensorOperator to nomalize the input image in preprocessing.
     */
    protected abstract TensorOperator getPreprocessNormalizeOp();

    /**
     * Gets the TensorOperator to dequantize the output probability in post processing.
     *
     * <p>For quantized model, we need de-quantize the prediction with NormalizeOp (as they are all
     * essentially linear transformation). For float model, de-quantize is not required. But to
     * uniform the API, de-quantize is added to float model too. Mean and std are set to 0.0f and
     * 1.0f, respectively.
     */
    protected abstract TensorOperator getPostprocessNormalizeOp();
}
